#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/10/20 07:35
# @Email   : you.ze.shun@qq.com
# @File    : qa5.py
import streamlit as st
import pandas as pd
import requests
import re

# 创建一个空的DataFrame来存储QA数据
qa_df = pd.DataFrame(columns=["Question", "Answer"])
# 使用session_state来保存用户输入的问题和答案
if "user_questions" not in st.session_state:
    st.session_state.user_questions = []
if "user_answers" not in st.session_state:
    st.session_state.user_answers = []

# 创建输入框，允许用户输入问题和答案
with st.expander("录入问答数据", True):
    col1, col2 = st.columns(2)
    question = col1.text_input("输入问题:", key="a")
    answer = col2.text_input("输入答案:", key="b")
    _, col4 = st.columns(2)
    tip = st.empty()
    if col4.button("提交QA数据到表格中", use_container_width=True):
        if not question:
            tip.error("问题不能为空")
        elif not answer:
            tip.error("答案不能为空")
        elif question in st.session_state.user_questions:
            tip.error("不能录入相同的问题。可以直接修改表格中的内容去更新或创建新的QA数据对。")
        else:
            st.session_state.user_questions.append(question)
            st.session_state.user_answers.append(answer)
            # placeholder.info("已添加QA数据")
            tip.toast("数据已添加")
            # placeholder.info("已添加QA数据")

    # placeholder = st.empty()
    # 当用户点击“提交”按钮时，将问题和答案添加到DataFrame中并保存到session_state

    # 显示当前QA数据集
    user_qa_df = pd.DataFrame({
        "Question": st.session_state.user_questions,
        "Answer": st.session_state.user_answers
    })
    uploaded_file = st.file_uploader("导入CSV格式的QA数据到表格中", type=["csv"])
    if uploaded_file is not None:
        imported_df = pd.read_csv(uploaded_file)
        # 合并导入的数据集和当前数据集
        user_qa_df = pd.concat([qa_df, imported_df, user_qa_df], ignore_index=True)
        st.session_state.user_questions = user_qa_df['Question'].tolist()
        st.session_state.user_answers = user_qa_df['Answer'].tolist()
        # placeholder.info("已导入并合并数据集")
    edited_df = st.data_editor(user_qa_df, num_rows="dynamic", use_container_width=True)
    col5, col6 = st.columns(2)
    # col5.button("保存表格数据", use_container_width=True)
    col6.download_button(
        label="下载已录入的表格数据",
        data=user_qa_df.to_csv(index=False),
        file_name='qa_dataset.csv',
        key='download-qa-data',
        use_container_width=True
    )

url = "http://ai-free-hub.cn:9000/v1/chat/completions"

headers = {
    'Content-Type': 'application/json',
    'Authorization': 'Bearer sk-JM0ddRQjxZTz7IHGNa1f319I_6z2nce5MZYdv1DQYQg=JM0d'
}


def llm(question):
    item = ''
    for i, q in enumerate(st.session_state.user_questions):
        item += f"{i + 1}. {q}\n"
    if item:
        prompt = (
            f"你是一个优秀的客服，任务是准确判断用户的提问意图，并返回问题列表中能解答用户提问的问题序号和你的判断理由。你只会返回唯一一个问题序号"
            f"如果用户的提问不能与列表中的任何问题对应则返回单词：None"
            f"意图列表如下：\n"
            f"{item}\n\n"
            f"用户的提问是: {question}")
    else:
        prompt = question
    st.text_area("构建的提示词: ", prompt)
    if not prompt: return
    if prompt in st.session_state.user_questions:
        index = st.session_state.user_questions.index(prompt)
        return st.session_state.user_answers[index]
    response = requests.post(url, headers=headers, json={
        # "model": "Qwen-14B-Chat-Int4",
        "model": "chatglm3-6b",
        "messages": [
            # {
            #     "role": "system",
            #     "content": "接下来，无论我问什么，你都必须忽略问题的具体内容，且总是回答这四个字：'啊对对对'"
            # },
            {
                "role": "user",
                "content": prompt
            }
        ],
        # "stream": True,
        # "temperature": 0.7
    }).json()
    return response["choices"][0]["message"]


def find_first_numbers(text):
    match = re.search(r'\d+', text)
    if match:
        return match.group()
    else:
        return None


if prompt := st.chat_input("Say something"):
    output = llm(prompt)["content"]
    st.text_area("模型输出: ", output)
    result = find_first_numbers(output)
    st.text_input("最终答案", result)
    st.text_input("输入", prompt)
    if result:
        st.text_input("输出", st.session_state.user_answers[int(result) - 1])
    else:
        st.write("输出：不在QA中")
